#%%
import gym
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.special import softmax
import numpy as np
import os

import matplotlib.pyplot as plt
from tqdm import tqdm
import GridWorld_v6
import draw
from collections import defaultdict, namedtuple

import argparse

import ppo
#%%
PPO_Transition = namedtuple('PPO_Transition', ['state', 'action', 'a_prob', 'reward', 'next_state','G'])# 定义PPO转换元组，用于存储经验
#%%
parser = argparse.ArgumentParser()
# * env
"""
未提及的均为0

targetAreaScore = 1
work

forbiddenAreaScore = -0.01
targetAreaScore = 1
work

forbiddenAreaScore = -0.01
hitWallScore = -0.01
targetAreaScore = 1
did't work

forbiddenAreaScore = -0.01
hitWallScore = -0.01
targetAreaScore = 1
entropy_coef = 0.01
work

"""



parser.add_argument("--forbiddenAreaScore", type=float, default = -0.01)  # 定义禁区得分参数
parser.add_argument("--targetAreaScore", type=float, default = 1)  # 定义目标区域得分参数
parser.add_argument("--hitWallScore", type=float, default = -0.01)  # 定义撞墙得分参数
parser.add_argument("--moveScore", type=float, default = 0)  # 定义移动得分参数
parser.add_argument("--desc", type=list, default = [".....",".##..","..#..",".#T#.",".#..."])  # 定义环境描述参数

# no need to change
parser.add_argument("--seed", type=int, default=0)  # 定义随机种子参数
parser.add_argument("--device", type=str, default="cpu")  # 定义设备参数
parser.add_argument("--action_space", type=int, default=4)  # 定义动作空间参数，动作空间是4个
parser.add_argument("--start_state", type=int, default=0)  # 定义起始状态参数，从十号点出发

# train
parser.add_argument("--num_episodes", type=int, default=20000)  # 定义训练回合数参数，从十号点出发

# * PPO hyperparameters
# PPO超参数
parser.add_argument("--gamma", type=float, default=0.99)  # 定义折扣因子参数
parser.add_argument("--batch_size", type=int, default=512)  # 定义批量大小参数
parser.add_argument("--num_update_per_iter", type=int, default=5)  # 定义每轨迹更新次数参数，每一个trajectory更新10次
parser.add_argument("--actor_lr", type=float, default=0.0003)  # 定义演员网络学习率参数
parser.add_argument("--critic_lr", type=float, default=0.001)  # 定义评论家网络学习率参数
parser.add_argument("--entropy_coef", type=float, default=0.01)  # 定义熵系数参数
parser.add_argument("--hidden_dim", type=int, default=64)  # 定义隐藏层维度参数
parser.add_argument("--clip_param", type=float, default=0.2)  # 定义裁剪参数
parser.add_argument("--max_grad_norm", type=float, default=5.0)  # 定义最大梯度范数参数
parser.add_argument("--trajectory_len", type=float, default=200)  # 定义轨迹长度参数

# * eval related
# 测试相关
parser.add_argument("--num_train_ep", type=int, default=10)  # 定义训练回合数参数
parser.add_argument("--num_eval_ep", type=int, default=5)  # 定义评估回合数参数
parser.add_argument("--ckpt_freq", type=int, default=100)  # 定义检查点频率参数
parser.add_argument("--eval_freq", type=int, default=100)  # 定义评估频率参数
parser.add_argument("--eval_runs", type=int, default=2)  # 定义评估运行次数参数

args = parser.parse_args(args=[])  # 解析命令行参数
#%%
def toOneHot(state):  # 定义状态转换为独热编码的函数
    stateid = state[0]*5 + state[1]  # 计算状态ID
    return np.eye(25)[stateid]  # 返回独热编码
def toCoordinate(oneHotState):  # 定义独热编码转换为坐标的函数
    stateid = np.argmax(oneHotState)  # 获取状态ID
    return [stateid//5, stateid%5]  # 返回坐标
#%%
def set_seed(seed):  # 定义设置随机种子的函数
    # do not fix the seed of the envs
    torch.set_num_threads(1)  # 设置PyTorch线程数
    torch.backends.cudnn.benchmark = False  # 禁用cudnn基准测试
    torch.backends.cudnn.deterministic = True  # 启用cudnn确定性
    
    torch.manual_seed(seed)  # 设置PyTorch随机种子
    torch.cuda.manual_seed_all(seed)  # 设置CUDA随机种子
    np.random.seed(seed)  # 设置NumPy随机种子

set_seed(args.seed)  # 设置随机种子
#%%
env = GridWorld_v6.GridWorld_v6(initState=0,  # 创建网格世界环境
                                moveScore=args.moveScore,  # 设置移动得分
                                action_space=args.action_space,  # 设置动作空间
                                forbiddenAreaScore=args.forbiddenAreaScore,  # 设置禁区得分
                                score=args.targetAreaScore,  # 设置目标区域得分
                                hitWallScore = args.hitWallScore,  # 设置撞墙得分
                                desc = args.desc,  # 设置环境描述
                                enterForbiddenArea=False)  # 设置是否可进入禁区

#%%
env.show()
#%%
env.reset() #重置环境
#%%
env.step(1) # 执行动作1
#%%
agent = ppo.PPO(state_dim = 25,  # 创建PPO代理
        hidden_dim = args.hidden_dim,  # 设置隐藏层维度
        action_dim = 4,  # 设置动作维度
        device = args.device,  # 设置设备
        num_steps = args.trajectory_len,  # 设置步数
        batch_size = args.batch_size,  # 设置批量大小
        actor_lr = args.actor_lr,  # 设置演员网络学习率
        critic_lr = args.critic_lr,  # 设置评论家网络学习率
        entropy_coef = args.entropy_coef,  # 设置熵系数
        gamma = args.gamma,  # 设置折扣因子
        num_update_per_iter = args.num_update_per_iter,  # 设置每轨迹更新次数
        clip_param = args.clip_param,  # 设置裁剪参数
        max_grad_norm = args.max_grad_norm)  # 设置最大梯度范数
#%%
agent.actor_net
#%%
env.show()
#%%
state = env.reset()
toOneHot(state)
#%%
np.array([toOneHot(state)], dtype=int)
#%%
agent.select_action(np.array([toOneHot(state)], dtype=int))
#%%
def collect_trajectory(env, agent, args, reward_list, test=True):  # 定义收集轨迹的函数
    state = env.reset()  # 重置环境并获取初始状态
    done = False  # 初始化完成标志
    cnt = 0  # 初始化计数器
    
    trans_list = []  # 初始化转换列表

    for i in range(args.trajectory_len):  # 循环轨迹长度次
        cnt += 1  # 计数器加1
        _, action, prob, _ = agent.select_action(np.array([toOneHot(state)], dtype=int))  # 选择动作
        action = action[0]  # 获取动作
        next_state, reward, done, _ = env.step(action)  # 执行动作并获取下一状态、奖励等
        trans = state, action, prob, reward, next_state  # 创建转换元组
        trans_list.append(trans)  # 将转换添加到列表
        state = next_state  # 更新状态
        reward_list.append(reward)  # 将奖励添加到列表
        if cnt > args.trajectory_len or done:  # 如果达到轨迹长度或完成
            break  # 跳出循环

    return_to_go = 0  # 初始化未来回报
    for trans in reversed(trans_list):  # 反向遍历转换列表,从目的地到起点
        state, action, prob, reward, next_state = trans  # 解包转换元组
        return_to_go = reward + args.gamma * return_to_go  # 计算未来回报G
        if not test:  # 如果不是测试
            agent.store_transition(PPO_Transition(toOneHot(state), action, prob, reward, toOneHot(next_state), return_to_go))  # 存储转换，把行为路径存储到buffer中

    return return_to_go  # 返回未来回报
#%%
from collections import Counter
#%%
results = []  # 初始化结果列表

eval_frequent = 50  # 设置评估频率

import copy  # 导入复制模块
pre_actor = copy.deepcopy(agent.actor_net.state_dict())  # 复制演员网络参数

# print(pre_actor)

for _ in tqdm(range(args.num_episodes)):  # 循环训练回合数次
    while(agent.buffer_size < agent.batch_size):  # 当缓冲区大小小于批量大小时
        reward_list = []  # 初始化奖励列表
        return_to_go = collect_trajectory(env, agent, args, reward_list , test=False)  # 收集轨迹
        # print(reward_list, return_to_go)
    
    # with torch.no_grad():
    # agent.actor_net.load_state_dict(pre_actor)
    # print(agent.buffer_size)
    
    if _ % eval_frequent==0:  # 如果是评估频率的倍数
        input =  torch.tensor([toOneHot((i,j)) for i in range(5) for j in range(5)],dtype=torch.float32)  # 创建输入张量
        pre_frame_probabilities = agent.actor_net(input).detach().numpy()  # 获取动作概率
        state_values = agent.critic_net(input).detach().numpy().reshape(5,5)  # 获取状态值

        # 绘制state_value
        p = pre_frame_probabilities.argmax(1)  # 获取最大概率的动作
        draw.draw(state_values, p)  # 绘制状态值和策略
        
        for i in range(5):  # 循环行
            for j in range(5):  # 循环列
                print(state_values[i][j], end="\t")  # 打印状态值
            print()  # 换行
        print('buffer_size',agent.buffer_size)  # 打印缓冲区大小
        transition_list = [(toCoordinate(t.state), t.action, t.reward, toCoordinate(t.next_state)) for t in list(agent.buffer)]  # 创建转换列表
        print(len(transition_list))  # 打印转换列表长度
    
    actor_loss, critic_loss, entropy = agent.update()  # 更新代理

    # agent.actor_net.load_state_dict(pre_actor)
    
    if _ % eval_frequent==0:  # 如果是评估频率的倍数
        input =  torch.tensor([toOneHot((i,j)) for i in range(5) for j in range(5)],dtype=torch.float32)  # 创建输入张量
        now_frame_probabilities = agent.actor_net(input).detach().numpy()  # 获取当前动作概率
        print(len(transition_list))  # 打印转换列表长度
        draw.plot_policy(pre_frame_probabilities, now_frame_probabilities,transition_list=transition_list)  # 绘制策略变化
    
        return_list = []  # 初始化回报列表
        for _ in range(5):  # 循环5次
            reward_list = []  # 初始化奖励列表
            return_to_go = collect_trajectory(env, agent, args, reward_list, test=True)  # 收集测试轨迹
            return_list.append(return_to_go)  # 将回报添加到列表
        result = sum(return_list) / len(return_list)  # 计算平均回报
        
        
#%%
import matplotlib.pyplot as plt

# 提供的数据列表
data = results
# 画折线图
plt.figure(figsize=(12, 6))
plt.plot(data, marker='o')
plt.title('Data Line Plot')
plt.xlabel('Index')
plt.ylabel('Value')
plt.grid(True)
plt.show()